from argparse import ArgumentParser
from ast import literal_eval
from time import time
from numpy.fft import fft2, fftshift, ifft2, ifftshift
from project_utils.image_io import *

import matplotlib.pyplot as plt
import numpy as np
import os.path as osp


def imnoise(gray_image, noise_type="gaussian", *args):
    if noise_type == "gaussian":
        if len(args) == 0:
            variance = 0.01
            mean = 0
        elif len(args) == 1:
            variance = 0.01
            mean = args[0]
        elif len(args) == 2:
            mean = args[0]
            variance = args[1]
        else:
            raise ValueError("Invalid number of arguments for 'gaussian' noise type")

        noise = np.random.normal(mean, np.sqrt(variance), gray_image.shape)
        ret_img = gray_image + noise

    elif noise_type == "poisson":
        ret_img = np.random.poisson(gray_image * 255) / 255.0

    elif noise_type == "salt & pepper":
        if len(args) == 0:
            density = 0.05
        elif len(args) == 1:
            density = args[0]
        else:
            raise ValueError(
                "Invalid number of arguments for 'salt & pepper' noise type"
            )

        noisy_pixels = np.random.choice(
            [-1, 1, 0], size=gray_image.shape, p=[density / 2, density / 2, 1 - density]
        )
        ret_img = gray_image
        # Add pepper noise
        ret_img[noisy_pixels == -1] = 0.0
        # Add salt noise
        ret_img[noisy_pixels == 1] = 1.0

    elif noise_type == "speckle":
        if len(args) == 0:
            variance = 0.05
        elif len(args) == 1:
            variance = args[0]
        else:
            raise ValueError("Invalid number of arguments for 'speckle' noise type")
        ret_img = gray_image + gray_image * np.random.normal(
            0, np.sqrt(variance), gray_image.shape
        )

    else:
        raise ValueError("Invalid noise type")

    return ret_img


def plot_image_on_ax(ax, image, title):
    """Plot image on the given axis."""

    plt.subplot(ax.get_subplotspec())
    plt.title(title, fontsize=8)
    plt.imshow(image, cmap="gray")
    plt.xticks([])
    plt.yticks([])


def unit_test_problem_01(image_path):
    """Imitate the imnoise, imnoise2, or imnoise3 functions from the textbook
    and provide Python implementations for adding various types of noise to images.
    """

    gray_image = load_image_as_grayscale(image_path, "rgb")
    gray_image = gray_image / 255.0

    fig, axes = plt.subplots(nrows=2, ncols=4, figsize=(12, 6))

    plot_image_on_ax(axes[0, 0], gray_image, "Raw image")

    noisy_gray_image = imnoise(gray_image, "gaussian")
    plot_image_on_ax(axes[0, 1], noisy_gray_image, "Gaussian Noise")

    noisy_gaussian_mean = imnoise(gray_image, "gaussian", 0.5)
    plot_image_on_ax(axes[0, 2], noisy_gaussian_mean, "Gaussian Noise (Custom Mean)")

    noisy_gaussian_mean_var = imnoise(gray_image, "gaussian", 0.5, 0.1)
    plot_image_on_ax(
        axes[0, 3], noisy_gaussian_mean_var, "Gaussian Noise (Custom Mean & Variance)"
    )

    noisy_poisson = imnoise(gray_image, "poisson")
    plot_image_on_ax(axes[1, 0], noisy_poisson, "Poisson Noise")

    noisy_salt_pepper = imnoise(gray_image, "salt & pepper")
    plot_image_on_ax(axes[1, 1], noisy_salt_pepper, "Salt & Pepper Noise")

    noisy_salt_pepper_density = imnoise(gray_image, "salt & pepper", 0.4)
    plot_image_on_ax(
        axes[1, 2], noisy_salt_pepper_density, "Salt & Pepper Noise (Custom Density)"
    )

    noisy_speckle = imnoise(gray_image, "speckle")
    plot_image_on_ax(axes[1, 3], noisy_speckle, "Speckle Noise")

    plt.show()


def mean_filter(signal, kernel_size):
    filtered_signal = np.zeros_like(signal)
    padding_left = int(kernel_size / 2)

    if kernel_size % 2 == 0:
        padding_right = int(kernel_size / 2) - 1
    else:
        padding_right = int(kernel_size / 2)

    for i in range(padding_left, len(signal) - padding_right):
        filtered_signal[i] = np.mean(signal[i - padding_left : i + padding_right])

    return filtered_signal


def median_filter(signal, kernel_size):
    filtered_signal = np.zeros_like(signal)
    padding_left = int(kernel_size / 2)

    if kernel_size % 2 == 0:
        padding_right = int(kernel_size / 2) - 1
    else:
        padding_right = int(kernel_size / 2)

    for i in range(padding_left, len(signal) - padding_right):
        filtered_signal[i] = np.median(signal[i - padding_left : i + padding_right])

    return filtered_signal


def unit_test_problem_02():
    """For non-periodic independent and identically distributed noise, design denoising
    functions based on various mean filters and order statistics filters. Provide test
    code to implement visual outputs for noise addition and denoising effects.
    """

    np.random.seed(0)
    N = 500
    x = np.linspace(0, 4 * np.pi, N)
    y = np.sin(x)

    noise = np.random.normal(0, 0.5, N)
    y_noisy = y + noise

    y_mean = mean_filter(y_noisy, 16)
    y_median = median_filter(y_noisy, 16)

    plt.figure(figsize=(12, 6))
    plt.plot(x, y, label="Original")
    plt.plot(x, y_noisy, label="Noisy")
    plt.plot(x, y_mean, label="Mean filter")
    plt.plot(x, y_median, label="Median filter")
    plt.legend()
    plt.show()


def unit_test_problem_03(image_path):
    img = load_image_as_grayscale(image_path, "rgb")
    height, width = img.shape

    # 指定最大滤波器大小
    max_filter_size = 11
    Median_Filter_Img = np.zeros_like(img)
    AdaMed_Filter_Img = np.zeros_like(img)

    # 边缘填充
    padding = int((max_filter_size - 1) / 2)
    image_padding = np.pad(
        img.copy(),
        (
            (padding, max_filter_size - padding - 1),
            (padding, max_filter_size - padding - 1),
        ),
        mode="edge",
    )

    # 计算常规中值滤波器的运行时间
    start_time = time()

    for x in range(padding, padding + height):
        for y in range(padding, padding + width):
            # 初始滤波器大小为3x3
            filter_size = 3

            # 常规中值滤波
            half_size = int(filter_size / 2)
            Padding = image_padding[
                x - half_size : x + half_size + 1, y - half_size : y + half_size + 1
            ]
            # 邻域 Sxy
            Median_Filter_Img[x - padding, y - padding] = np.median(Padding)

    end_time = time()
    conventional_time = end_time - start_time
    print(f"常规中值滤波器运行时间: {conventional_time:.2f}s")

    # 计算自适应中值滤波器的运行时间
    start_time = time()

    for x in range(padding, padding + height):
        for y in range(padding, padding + width):
            # 初始滤波器大小为3x3
            filter_size = 3

            # 邻域 Sxy 中心像素
            Zxy = img[x - padding, y - padding]

            while filter_size < max_filter_size:
                # 自适应中值滤波器
                half_size = int(filter_size / 2)
                Padding = image_padding[
                    x - half_size : x + half_size + 1, y - half_size : y + half_size + 1
                ]
                # 邻域 Sxy
                Zmin = np.min(Padding)
                Zmax = np.max(Padding)
                Zmed = np.median(Padding)

                if Zmed > Zmin and Zmax > Zmed:
                    if Zxy > Zmin and Zmax > Zxy:
                        AdaMed_Filter_Img[x - padding, y - padding] = Zxy
                    else:
                        AdaMed_Filter_Img[x - padding, y - padding] = Zmed
                    break

                else:
                    filter_size = filter_size + 2

            else:
                AdaMed_Filter_Img[x - padding, y - padding] = Zxy

    end_time = time()
    adaptive_time = end_time - start_time
    print(f"自适应中值滤波器运行时间: {adaptive_time:.2f}s")

    # 保存结果
    image_folder = osp.dirname(image_path)
    save_grayscale_image(
        osp.join(image_folder, "Common_Median_Filter.jpg"), Median_Filter_Img
    )
    save_grayscale_image(
        osp.join(image_folder, "Adaptive_Median_Filter.jpg"), AdaMed_Filter_Img
    )


def plot_spectrum(image_path):
    image = load_image_as_grayscale(image_path, "rgb")
    height, width = image.shape

    # 对信号进行二维傅里叶变换并进行中心化
    spectrum = fftshift(fft2(image))
    magnitude_spectrum = np.abs(spectrum)

    plt.imshow(
        np.log(1 + magnitude_spectrum),
        cmap="gray",
        extent=[-width / 2, width / 2, -height / 2, height / 2],
    )

    plt.colorbar()
    plt.show()


# 在频域上实现带阻滤波器
def bandstop_filter_2d(image, lowcut, highcut):
    rows, cols = image.shape
    y, x = np.ogrid[-rows / 2 : rows / 2, -cols / 2 : cols / 2]
    mask = np.zeros((rows, cols))

    # 创建带阻滤波器
    mask[(x**2 + y**2) < lowcut**2] = 1
    mask[(x**2 + y**2) > highcut**2] = 1

    # 对原始图像进行傅里叶变换
    spectrum = fftshift(fft2(image))

    # 将图像和滤波器相乘
    filtered_spectrum = spectrum * mask

    # 对结果进行反变换
    filtered_image = np.abs(ifft2(ifftshift(filtered_spectrum)))

    return filtered_image


# 在频域上实现陷波滤波器
def notch_filter_2d(image, u0, u1, D0, n):
    rows, cols = image.shape
    y, x = np.ogrid[-rows / 2 : rows / 2, -cols / 2 : cols / 2]

    D1 = np.sqrt((x - u0) ** 2 + (y - u1) ** 2)
    D2 = np.sqrt((x + u0) ** 2 + (y + u1) ** 2)

    # n阶巴特沃斯带阻滤波器
    H = 1 / (1 + ((D0 * D0) / (D1 * D2 + 1e-8)) ** n)

    # 理想带阻滤波器
    # H = np.ones((rows, cols))
    # H[D1 <= D0] = 0
    # H[D2 <= D0] = 0

    # 对原始图像进行傅里叶变换
    spectrum = fftshift(fft2(image))

    # 将图像和滤波器相乘
    filtered_spectrum = spectrum * H

    # 对结果进行反变换
    filtered_image = np.abs(ifft2(ifftshift(filtered_spectrum)))

    return filtered_image


# 在频域上实现最优陷波滤波器
def optimal_notch_filter_2d(image, u, D0, n):
    rows, cols = image.shape
    y, x = np.ogrid[-rows / 2 : rows / 2, -cols / 2 : cols / 2]

    H = np.ones((rows, cols))
    for [u0, u1] in u:
        D1 = np.sqrt((x - u0) ** 2 + (y - u1) ** 2)
        D2 = np.sqrt((x + u0) ** 2 + (y + u1) ** 2)

        # n阶巴特沃斯带阻滤波器
        H = H / (1 + ((D0 * D0) / (D1 * D2 + 1e-8)) ** n)

        # 理想带阻滤波器
        # H[D1 <= D0] = 0
        # H[D2 <= D0] = 0

    # 对原始图像进行傅里叶变换
    spectrum = fftshift(fft2(image))

    # 将图像和滤波器相乘
    filtered_spectrum = spectrum * H

    # 对结果进行反变换
    filtered_image = np.abs(ifft2(ifftshift(filtered_spectrum)))

    raw_noise = image - filtered_image
    term_a = np.mean(image * raw_noise)
    term_b = np.mean(image) * np.mean(raw_noise)
    term_c = np.mean(raw_noise * raw_noise)
    term_d = np.mean(raw_noise) * np.mean(raw_noise)
    w = (term_a - term_b) / (term_c - term_d)

    filtered_image = image - w * raw_noise

    return filtered_image


def unit_test_problem_04_1(image_path, cmdargs):
    lowcut, highcut = cmdargs.lowcut, cmdargs.highcut

    original_image = load_image_as_grayscale(image_path, "rgb")
    filtered_image = bandstop_filter_2d(original_image, lowcut, highcut)

    fig, axes = plt.subplots(1, 2, figsize=(12, 6))

    plot_image_on_ax(axes[0], original_image, "Original")
    plot_image_on_ax(axes[1], filtered_image, "Bandreject")

    plt.show()


def unit_test_problem_04_2(image_path, cmdargs):
    u0, u1 = cmdargs.notch_u0, cmdargs.notch_u1
    D0, n = cmdargs.notch_D0, cmdargs.notch_n

    original_image = load_image_as_grayscale(image_path, "rgb")
    filtered_image = notch_filter_2d(original_image, u0, u1, D0, n)

    fig, axes = plt.subplots(1, 2, figsize=(12, 6))

    plot_image_on_ax(axes[0], original_image, "Original")
    plot_image_on_ax(axes[1], filtered_image, "Notch")

    plt.show()


def unit_test_problem_04_3(image_path, cmdargs):
    u = cmdargs.optimal_notch_centers
    D0, n = cmdargs.notch_D0, cmdargs.notch_n

    original_image = load_image_as_grayscale(image_path, "rgb")
    filtered_image = optimal_notch_filter_2d(original_image, u, D0, n)

    fig, axes = plt.subplots(1, 2, figsize=(12, 6))

    plot_image_on_ax(axes[0], original_image, "Original")
    plot_image_on_ax(axes[1], filtered_image, "Optimal Notch")

    plt.show()


def python_expr(literal: str, default=None):
    """Type definition for `argparse.ArgumentParser`
    that parses an integer, a tuple or a list.
    """

    try:
        py_value = literal_eval(literal.strip())
    except Exception as original_exception:
        if default is not None:
            py_value = default
        else:
            raise original_exception

    return py_value


def main_procedure(cmdargs):
    if cmdargs.mode == "add-noises":
        image_path = osp.abspath(cmdargs.image)
        unit_test_problem_01(image_path)

    if cmdargs.mode == "visualize-denoising":
        unit_test_problem_02()

    if cmdargs.mode == "adaptive-median":
        image_path = osp.abspath(cmdargs.image)
        unit_test_problem_03(image_path)

    if cmdargs.mode == "spectrum":
        image_path = osp.abspath(cmdargs.image)
        plot_spectrum(image_path)

    if cmdargs.mode == "bandreject":
        image_path = osp.abspath(cmdargs.image)
        unit_test_problem_04_1(image_path, cmdargs)

    if cmdargs.mode == "notch":
        image_path = osp.abspath(cmdargs.image)
        unit_test_problem_04_2(image_path, cmdargs)

    if cmdargs.mode == "optimal-notch":
        image_path = osp.abspath(cmdargs.image)
        unit_test_problem_04_3(image_path, cmdargs)


if __name__ == "__main__":
    parser = ArgumentParser(description="Denoising Algorithms.")

    parser.add_argument(
        "mode",
        choices=[
            "add-noises",
            "visualize-denoising",
            "adaptive-median",
            "spectrum",
            "bandreject",
            "notch",
            "optimal-notch",
        ],
        help="Select one of the unit tests.",
    )

    parser.add_argument("--image", type=str, help="Filepath of the test image.")

    parser.add_argument(
        "--lowcut", type=float, help="Cut-off frequency (Low) in bandreject filtering."
    )
    parser.add_argument(
        "--highcut",
        type=float,
        help="Cut-off frequency (High) in bandreject filtering.",
    )

    parser.add_argument("--notch_u0", type=float, help="Reject center (X).")
    parser.add_argument("--notch_u1", type=float, help="Reject center (Y).")
    parser.add_argument("--notch_D0", type=float, help="Radius of band center.")
    parser.add_argument("--notch_n", type=float, help="Order of notch filter.")

    parser.add_argument(
        "--optimal_notch_centers", type=python_expr, help="Reject centers."
    )

    main_procedure(parser.parse_args())
